package gov.fnal.xpath2ldap;

import org.w3c.dom.*;
import org.xml.sax.SAXException;
import org.xml.sax.InputSource;

import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.dom.DOMResult;
import javax.xml.transform.URIResolver;
import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.TransformerException;
import javax.xml.transform.Source;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.ParserConfigurationException;
import javax.xml.parsers.DocumentBuilder;
import java.io.PrintWriter;
import java.io.IOException;
import java.io.FileOutputStream;
import java.util.*;
import java.net.URL;

/**
 * User: tdh
 * Date: Mar 4, 2009
 * Time: 2:44:43 PM
 */
public class DSMLTransform implements URIResolver {

    TransformerFactory tFactory;
    DocumentBuilder dBuilder;
    Transformer transformer;
    Element first_directory_schema_element=null;
    Transformer merge_transformer;
    LinkedHashMap<String, Document> mergeDocuments = new LinkedHashMap<String, Document>();
    XMLAdapter xmler;

    public DSMLTransform(XMLAdapter xmler) {
        this.xmler = xmler;
        tFactory = TransformerFactory.newInstance();

        if(tFactory.getFeature(DOMSource.FEATURE) && tFactory.getFeature(DOMResult.FEATURE))
        {
            try {
                //Instantiate a DocumentBuilderFactory.
                System.setProperty("javax.xml.parsers.XmlDocumentParserFactory", "org.apache.xerces.jaxp.DocumentBuilderFactoryImpl");
                DocumentBuilderFactory dFactory = DocumentBuilderFactory.newInstance();

                // And setNamespaceAware, which is required when parsing xsl files
                dFactory.setNamespaceAware(true);

                //Use the DocumentBuilderFactory to create a DocumentBuilder.
                dBuilder = dFactory.newDocumentBuilder();

                //Use the DocumentBuilder to parse the XSL stylesheet.
                //todo: remove hard-coded xsl transform name
                URL dsml2xml = DSMLTransform.class.getResource("/dsml2xml.xsl");
                Document xslDoc = dBuilder.parse(dsml2xml.openStream());

                // Use the DOM Document to define a DOMSource object.
                DOMSource xslDomSource = new DOMSource(xslDoc);

                // Set the systemId: note this is actually a URL, not a local filename
                xslDomSource.setSystemId("dsml2xml.xsl");

                // Process the stylesheet DOMSource and generate a Transformer.
                transformer = tFactory.newTransformer(xslDomSource);

                //Get the schema to add to the source.
                URL glueldapschema = DSMLTransform.class.getResource("/glue-1.3-ldap.xml");
                Document schemaDoc = dBuilder.parse(glueldapschema.openStream());
                //todo: get schema from ldapsearch, cache as file

                // Get the directory schema element, to add to the search results document
                NodeList schemas = schemaDoc.getElementsByTagName("directory-schema");
                if(schemas.getLength()>=1) {
                    Node firstone = schemas.item(0);
                    if (firstone instanceof Element) {
                        first_directory_schema_element = (Element) firstone;
                    }
                }

            } catch (ParserConfigurationException pce) {
                System.err.println("ParserConfigurationException");
            } catch (IOException pce) {
                System.err.println("IOException");
            } catch (TransformerException pce) {
                System.err.println("TransformerException");
            } catch (SAXException pce) {
                System.err.println("SAXException");
            }
        }  else
        {
            System.err.println("DOM node processing not supported!");
        }
    }

    // Rewrite dsml results document, adding ancestor elements with parent-child structure
    // The dsml element should be the document root, but write for generality.
    public Document xformDSMLToStructuredXML(Pair<Document, LinkedHashMap<String, Stack<String>>> dsmlPair) {
        Document xmlforattr = dsmlPair.getFirst();
        LinkedHashMap<String, Stack<String>> ancestorMap = dsmlPair.getSecond();
        Element xmlforattr_dsml_element= null;
        NodeList dsml_elts = xmlforattr.getElementsByTagName("dsml");
        if(dsml_elts.getLength()>=1) {
            Node firstone = dsml_elts.item(0);
            if (firstone instanceof Element) {
                xmlforattr_dsml_element = (Element) firstone;
            }
        } else {
            System.err.println("Returned ldap search format was incorrect -no dsml element was found.");
            return xmlforattr;
        }
        // Will add created ancestor entries to this element
        Element xmlforattr_directory_entries=null;
        NodeList entries = xmlforattr_dsml_element.getElementsByTagName("directory-entries");
        if(entries.getLength()>=1) {
            Node firstone = entries.item(0);
            if (firstone instanceof Element) {
                xmlforattr_directory_entries = (Element) firstone;
            }
        } else {
            System.err.println("Returned ldap search format was incorrect -no directory-entries element was found.");
            return xmlforattr;
        }


        if(false) {
            try {
                FileOutputStream fos = new FileOutputStream("/home/tdh/tmp/fromsearch-5.forattr.xml");
                XMLAdapter.serialize(xmlforattr, fos);
                fos.flush();
                fos.close();
            } catch (Exception e) {
                System.err.println("Could not open output file or transfrom search result xml to a string");
            }
        }

        Document ancestorDoc  = dBuilder.getDOMImplementation().createDocument(null, "dsml", null);
        Element dsml_element = ancestorDoc.getDocumentElement();
        Element directory_entries = (Element) dsml_element.appendChild(ancestorDoc.createElement("directory-entries"));
        boolean use_ldap_aux = xmler.getCommandLine().hasOption(XPathSearchOptions.USE_LDAP_AUX);
        directory_entries.setAttribute("use_ldap_aux", Boolean.valueOf(use_ldap_aux).toString());


        LinkedHashMap<String, Element> dsmlEntries = new LinkedHashMap<String, Element>();
        // For each entry, place it under its parent element, creating ancestors if necesary.
        NodeList entrynodes = xmlforattr_directory_entries.getChildNodes();
        int numentries = entrynodes.getLength();
        for(int i = 0; i < numentries; i++) {
            Element dsml_entry = (Element)entrynodes.item(i);
            String entryname = dsml_entry.getNodeName();
            if(dsml_entry==null) {
                System.err.println(i + " " + "found null");
                continue;
            }
            String entry_dn = dsml_entry.getAttribute("dn");
            // Check whether there is already an entry for this dn
            Element entryElement = dsmlEntries.get(entry_dn);
            if (entryElement!=null) {
                // Make a hashmap of the latest attributes
                LinkedHashMap<String, Element> dsml_entry_attr_map = new LinkedHashMap<String, Element> ();
                List<Element> dsml_entry_attr_elts = getChildElementsByTagName(dsml_entry, "attr");
                for(Element attr_elt : dsml_entry_attr_elts) {
                        String attr_name = attr_elt.getAttributeNode("name").getValue();
                        dsml_entry_attr_map.put(attr_name, attr_elt);
                }
                // Update attribute elements in existing entry element
                LinkedHashSet<String> entryElement_attr_set = new LinkedHashSet<String> ();
                List<Element>  entryElement_attr_elts = getChildElementsByTagName(entryElement, "attr");
                for(Element attr_elt : entryElement_attr_elts) {
                    String attr_name = attr_elt.getAttributeNode("name").getValue();
                    Element dsml_entry_attr_elt = dsml_entry_attr_map.get(attr_name);
                    if(dsml_entry_attr_elt!=null) {
                        Node replacewith = ancestorDoc.importNode(dsml_entry_attr_elt, true);
                        entryElement.replaceChild(replacewith, attr_elt);
                        // Remove this attribute from the set, so that it is not also added later
                        dsml_entry_attr_map.remove(attr_name);
                    }
                }
                // Add attribute elements to existing entry element
                for(Element attr_elt : dsml_entry_attr_map.values()) {
                    Node addwith = ancestorDoc.importNode(attr_elt, true);
                    entryElement.appendChild(addwith);
                }
                // Entry has been merged.
            } else {
                // Acquire the classname
                String classname=null;
                // Find the objectclass child
                Element objectclass= null;
                List<Element> oc_elts = getChildElementsByTagName(dsml_entry, "objectclass");
                if(oc_elts.size()>=1) {
                    Node firstone = oc_elts.get(0);
                    if (firstone instanceof Element) {
                        objectclass = (Element) firstone;
                    }
                } else {
                    System.err.println("Returned ldap search format was incorrect -no objectclass element was found for " + dsml_entry);
                    continue;
                }
                // Find its oc-value children
                List<Element> ocval_elts = getChildElementsByTagName(objectclass, "oc-value");
                for(Element oc_value_elt : ocval_elts) {
                        classname = oc_value_elt.getTextContent();
                        if(classname==null) {
                            continue;
                        }
                        if (isStructuralObjectclass(classname)) {
                            break;
                        }
                }
                if(classname==null) {
                    System.err.println("Returned ldap search format was incorrect -no classname could be determined for " + dsml_entry);
                    continue;
                }
                // Find or create the ancestors as far back as the earliest common one
                Stack<String> ancestor_classnames = ancestorMap.get(classname);
                // Convenience stack for constructing structure later
                Stack<Element> ancestor_elements = new Stack<Element>();
                Element next_ancestor=null;
                Element earliest_ancestor=null;
                Iterator<String> ancestor_classnames_iterator = ancestor_classnames.iterator();
                String ancestor_dn = entry_dn;
                while (ancestor_dn.length()>0 && ancestor_classnames_iterator.hasNext()) {
                    // Look for next ancestor in dn
                    int comma_index = ancestor_dn.indexOf(',');
                    if(comma_index<0) break;
                    ancestor_dn = ancestor_dn.substring(comma_index+1);
                    // Check whether the parent dn is in the hashtable
                    next_ancestor = dsmlEntries.get(ancestor_dn);
                    // Break, found existing ancestor
                    if (next_ancestor!=null) break;
                    // Create the ancestor and then continue
                    next_ancestor = ancestorDoc.createElement("entry");
                    next_ancestor.setAttribute("dn", ancestor_dn);
                    Element occlasses_element = ancestorDoc.createElement("objectclass");
                    next_ancestor.appendChild(occlasses_element);
                    Element ocvalue_element = ancestorDoc.createElement("oc-value");
                    occlasses_element.appendChild(ocvalue_element);
                    String ancestor_name = ancestor_classnames_iterator.next();
                    Text ocvalue_text = ancestorDoc.createTextNode(ancestor_name);
                    ocvalue_element.appendChild(ocvalue_text);
                    dsmlEntries.put(ancestor_dn, next_ancestor);
                    earliest_ancestor = next_ancestor;
                    ancestor_elements.push(next_ancestor);
                }

                // Put together the "structured dsml" entries.
                Element addto_ancestor;
                // No already-created ancestors were found -must add earliest to directory-entries
                if(earliest_ancestor==next_ancestor) {
                    addto_ancestor = directory_entries;
                } else {
                    addto_ancestor = next_ancestor;
                }
                // This loop starts from the earliest common ancestor, going towards dsml_entry
                while(!ancestor_elements.empty()) {
                    next_ancestor = ancestor_elements.pop();
                    addto_ancestor.appendChild(next_ancestor);
                    addto_ancestor = next_ancestor;
                }
                // Add the entry to its closest ancestor
                Node imported_dsml_entry = ancestorDoc.importNode(dsml_entry, true);
                addto_ancestor.appendChild(imported_dsml_entry);
                // Place this entry in the hashtable
                dsmlEntries.put(entry_dn, (Element) imported_dsml_entry);
            }
        }

        if(use_ldap_aux) {
            // Fix the attributes.
            // If an attribute belongs to an auxiliary class, that class should be its parent.
            Map<String, Map<String, String>> true_parent_name_map = new LinkedHashMap<String, Map<String, String>>();
            for (Element entryElement : dsmlEntries.values()) {
                String entry_dn = entryElement.getAttribute("dn");
                List<Element> entryElement_attr_elts = getChildElementsByTagName(entryElement, "attr");
                int num_entryElement_attrs = entryElement_attr_elts.size();
                // No attributes to fix
                if(num_entryElement_attrs==0) {
                    continue;
                }

                // Get a list of the child classes, which are the oc-value elements
                List<Element> oc_elts = getChildElementsByTagName(entryElement, "objectclass");
                Element objectclass=null;
                if(oc_elts.size()>=1) {
                    Node firstone = oc_elts.get(0);
                    if (firstone instanceof Element) {
                        objectclass = (Element) firstone;
                    }
                } else {
                    System.err.println("DSML search result was incorrect -no objectclass element was found for " + entryElement);
                    continue;
                }

                // The first of these that is structural is the name of the objectclass
                List<Element> ocvalues_list = getChildElementsByTagName(objectclass, "oc-value");
                String objectclass_name=null;
                for(Element oc_value_elt : ocvalues_list) {
                    objectclass_name = oc_value_elt.getTextContent();
                    if(objectclass_name==null) {
                        continue;
                    }
                    if (isStructuralObjectclass(objectclass_name)) {
                        break;
                    }
                }

                LinkedHashMap<String, Element> aux_children = new LinkedHashMap<String, Element>();
                // Get map of attributes to true parent names for this class
                Map<String, String> attr_true_parent_map = true_parent_name_map.get(objectclass_name);
                if(attr_true_parent_map==null) {
                    attr_true_parent_map = new LinkedHashMap<String, String>();
                    true_parent_name_map.put(objectclass_name, attr_true_parent_map);
                }

                // Go through the attributes and move them if necesary
                for(Node next_node : entryElement_attr_elts) {
                    Element attr_elt = (Element) next_node;
                    String attr_name = attr_elt.getAttributeNode("name").getValue();
                    String true_parent_name = attr_true_parent_map.get(attr_name);
                    // Find an auxilary class that has this attribute
                    if(true_parent_name==null) {
                        // If there are none, this class is the parent of the attribute
                        true_parent_name = objectclass_name;
                        for (Element oc_value_elt : ocvalues_list) {
                            String classname = oc_value_elt.getTextContent();
                            if(classname==null) {
                                continue;
                            }
                            if (isAuxiliaryObjectclass(classname)) {
                                if(hasAtribute(classname, attr_name)) {
                                    true_parent_name = classname;
                                    break;
                                }
                            }
                        }
                        attr_true_parent_map.put(attr_name, true_parent_name);
                    }

                    // If such is found, move the attribte to that element
                    if(!true_parent_name.equals(objectclass_name)) {
                        // See if this auxilary child has already been added
                        Element aux_child = aux_children.get(true_parent_name);
                        if(aux_child==null) {
                            // No, create it
                            aux_child = ancestorDoc.createElement("entry");
                            aux_child.setAttribute("dn", true_parent_name + "=auxiliary," + entry_dn);
                            Element occlasses_element = ancestorDoc.createElement("objectclass");
                            Element ocvalue_element = ancestorDoc.createElement("oc-value");
                            Text ocvalue_text = ancestorDoc.createTextNode(true_parent_name);
                            ocvalue_element.appendChild(ocvalue_text);
                            occlasses_element.appendChild(ocvalue_element);
                            aux_child.appendChild(occlasses_element);
                            entryElement.appendChild(aux_child);
                            aux_children.put(true_parent_name, aux_child);
                        }
                        entryElement.removeChild(attr_elt);
                        aux_child.appendChild(attr_elt);
                    }
                }
            }
        }

        if(false)
            try {
                FileOutputStream fos = new FileOutputStream("/home/tdh/tmp/fromsearch-5.dsml.xml");
                XMLAdapter.serialize(ancestorDoc, fos);
                fos.flush();
                fos.close();
            } catch (Exception e) {
                System.err.println("Could not transfrom search result xml to a string");
            }

        // Transform the structured dsml document into xml with element names as per the ldap schema.
        // Insert the schema before the entries
        if(first_directory_schema_element==null) { //todo: allow xslt to access schema without the insertion.
            //todo: Extract only those parts of schema that are needed.
            System.err.println("No schema found.");
        }
        Node imported_schema_elt = ancestorDoc.importNode(first_directory_schema_element, true);
        if(first_directory_schema_element==null) {
            System.err.println("No entries found.");
        }
        dsml_element.insertBefore(imported_schema_elt, directory_entries);
        if(false)
            try {
                FileOutputStream fos = new FileOutputStream("/home/tdh/tmp/fromsearch-5.dsml.schema.xml");
                XMLAdapter.serialize(ancestorDoc, fos);
                fos.flush();
                fos.close();
            } catch (Exception e) {
                System.err.println("Could not transfrom search result xml to a string");
            }

        if(false) {
            InputSource source = new InputSource("/home/tdh/tmp/fromsearch-5.dsml.schema.xml"); //todo: need utf-8
            DocumentBuilderFactory dbf = null;
            DocumentBuilder db = null;
            try {
                if ( dbf == null ) {
                    dbf = DocumentBuilderFactory.newInstance();
                    dbf.setNamespaceAware( true );
                    dbf.setValidating( false );
                }

            db = dbf.newDocumentBuilder();
            ancestorDoc = db.parse( source );
                } catch ( Exception e ) {
                e.printStackTrace();
            }
            }

        // Transform the structured dsml document into xml with element names as per the ldap schema.
        DOMSource xmlDomSource = new DOMSource(ancestorDoc);
        xmlDomSource.setSystemId("ancestorDoc");
        DOMResult domResult = new DOMResult();
        try {
            // Use schema to convert from dsml
            transformer.transform(xmlDomSource, domResult); //todo: increase speed of dsml2xml transform
            //todo: try merging all dsml entries first again
        } catch (
                TransformerException pce) {
            System.err.println("TransformerException");
        }

        return (Document) domResult.getNode();
    }

    public Source resolve(String publicId, String systemId)
    {
        //System.out.println("resolving " + publicId + " and  " + systemId);
        if (true) {
            return new DOMSource(mergeDocuments.get(publicId));
        }

        return null;
    }

    public boolean isAuxiliaryObjectclass(String classname) {
        return xmler.isAuxiliaryObjectclass(classname);
    }

    public boolean isStructuralObjectclass(String classname) {
        return xmler.isStructuralObjectclass(classname);
    }

    public boolean hasAtribute(String classname, String attrname) {
        return xmler.hasAtribute(classname, attrname);
    }

    public static List<Element> getChildElementsByTagName(Element elt, String tagname) {
        List children = new ArrayList<Element>();
        // Find all children
        NodeList child_nodes = elt.getChildNodes();
        int num_children = child_nodes.getLength();
        // Check their names
        for(int i=0; i<num_children;i++) {
            Node child = child_nodes.item(i);
            if (child instanceof Element) {
                if(child.getNodeName().equals(tagname)) {
                children.add(child);
                }
            }
        }
        return children;
    }
}
